// Taken from candle-examples
use hf_hub::api::sync::ApiRepo;
use std::{collections::HashSet, fs, path::PathBuf};

/// Loads the safetensors files for a model from the hub based on a json index file.
pub fn hub_load_safetensors(repo: &ApiRepo, json_file: &str) -> candle_core::Result<Vec<PathBuf>> {
    let json_file = repo.get(json_file).map_err(candle_core::Error::wrap)?;
    let json_file = fs::File::open(json_file)?;
    let json: serde_json::Value =
        serde_json::from_reader(&json_file).map_err(candle_core::Error::wrap)?;
    let weight_map = match json.get("weight_map") {
        None => candle_core::bail!("no weight map in {json_file:?}"),
        Some(serde_json::Value::Object(map)) => map,
        Some(_) => candle_core::bail!("weight map in {json_file:?} is not a map"),
    };

    let mut safetensors_files = HashSet::new();
    for value in weight_map.values() {
        if let Some(file) = value.as_str() {
            safetensors_files.insert(file.to_string());
        }
    }

    let safetensors_files = safetensors_files
        .iter()
        .map(|v| repo.get(v).map_err(candle_core::Error::wrap))
        .collect::<candle_core::Result<Vec<_>>>()?;

    Ok(safetensors_files)
}
